#!/usr/bin/env python3
"""
file_agent.py
Capture openat/write syscalls to record filenames and volume of writes per PID.
Requires sudo.
Fixed version: Replaced ssize_t with int to avoid type recognition issues.
"""
from bcc import BPF
import os, json
from datetime import datetime, timezone

OUT_DIR = "/var/log/os_monitor"
os.makedirs(OUT_DIR, exist_ok=True)
OUT_FILE = os.path.join(OUT_DIR, "fileio.jsonl")

BPF_PROGRAM = r"""
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

struct data_t {
    u32 pid;
    u64 ts_ns;
    int event; // 1 openat, 2 write
    char fname[256];
    u64 count;  // Changed from ssize_t to u64 for compatibility
};

BPF_PERF_OUTPUT(events);

int trace_openat(struct pt_regs *ctx) {
    struct data_t data = {};
    data.ts_ns = bpf_ktime_get_ns();
    data.pid = bpf_get_current_pid_tgid() >> 32;
    data.event = 1;
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}

int trace_write_entry(struct pt_regs *ctx) {
    struct data_t data = {};
    data.ts_ns = bpf_ktime_get_ns();
    data.pid = bpf_get_current_pid_tgid() >> 32;
    data.event = 2;
    // For simplicity, we'll get the count from the return value in trace_write_return
    events.perf_submit(ctx, &data, sizeof(data));
    return 0;
}
"""

b = BPF(text=BPF_PROGRAM)

# 使用更兼容的tracepoint方式
try:
    # 使用tracepoint而不是kprobe，更稳定
    b.attach_tracepoint(tp="syscalls:sys_enter_write", fn_name="trace_write_entry")
except Exception as e:
    print(f"Warning: Could not attach to write tracepoint: {e}")
    # 回退到kprobe
    try:
        b.attach_kprobe(event="__x64_sys_write", fn_name="trace_write_entry")
    except Exception as e2:
        print(f"Warning: Could not attach to write kprobe: {e2}")

try:
    b.attach_tracepoint(tp="syscalls:sys_enter_openat", fn_name="trace_openat")
except Exception as e:
    print(f"Warning: Could not attach to openat tracepoint: {e}")

def write_record(record: dict):
    with open(OUT_FILE, "a") as f:
        f.write(json.dumps(record, default=str) + "\n")

def print_event(cpu, data, size):
    try:
        event = b["events"].event(data)
        r = {
            "pid": int(event.pid),
            "ts_ns": int(event.ts_ns),
            "event": int(event.event),
            "count": int(event.count) if event.event == 2 else 0,
            "ts": datetime.now(timezone.utc).isoformat()
        }
        write_record(r)
    except Exception as e:
        print(f"Error processing event: {e}")

b["events"].open_perf_buffer(print_event)
print("[+] file_agent running, writing to {}".format(OUT_FILE))

try:
    while True:
        b.perf_buffer_poll()
except KeyboardInterrupt:
    print("\nShutting down file_agent...")